import os
import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import threading
from loguru import logger
from apis.pc_apis import XHS_Apis
from xhs_utils.data_util import handle_note_info, download_note, save_to_xlsx

class Data_Spider():
    def __init__(self):
        self.xhs_apis = XHS_Apis()

    def spider_note(self, note_url: str, cookies_str: str, proxies=None):
        """
        爬取一个笔记的信息
        :param note_url:
        :param cookies_str:
        :return:
        """
        note_info = None
        try:
            success, msg, note_info = self.xhs_apis.get_note_info(note_url, cookies_str, proxies)
            if success:
                note_info = note_info['data']['items'][0]
                note_info['url'] = note_url
                note_info = handle_note_info(note_info)
        except Exception as e:
            success = False
            msg = e
        logger.info(f'爬取笔记信息 {note_url}: {success}, msg: {msg}')
        return success, msg, note_info

    def spider_some_note(self, notes: list, cookies_str: str, base_path: dict, save_choice: str, excel_name: str = '', proxies=None):
        """
        爬取一些笔记的信息
        :param notes:
        :param cookies_str:
        :param base_path:
        :return:
        """
        if (save_choice == 'all' or save_choice == 'excel') and excel_name == '':
            raise ValueError('excel_name 不能为空')
        note_list = []

        # 自动创建媒体文件夹
        if save_choice in ('all', 'media'):
            media_path = base_path['media']
            os.makedirs(media_path, exist_ok=True)
            for note_info in note_list:
                download_note(note_info, media_path)

        # 自动创建Excel文件夹并保存
        if save_choice in ('all', 'excel'):
            excel_path = base_path['excel']
            os.makedirs(excel_path, exist_ok=True)
            file_path = os.path.join(excel_path, f"{excel_name}.xlsx")
            save_to_xlsx(note_list, file_path)

        for note_url in notes:
            success, msg, note_info = self.spider_note(note_url, cookies_str, proxies)
            if note_info is not None and success:
                note_list.append(note_info)
        for note_info in note_list:
            if save_choice == 'all' or save_choice == 'media':
                download_note(note_info, base_path['media'])
        if save_choice == 'all' or save_choice == 'excel':
            file_path = os.path.abspath(os.path.join(base_path['excel'], f'{excel_name}.xlsx'))
            save_to_xlsx(note_list, file_path)


    def spider_user_all_note(self, user_url: str, cookies_str: str, base_path: dict, save_choice: str, excel_name: str = '', proxies=None):
        """
        爬取一个用户的所有笔记
        :param user_url:
        :param cookies_str:
        :param base_path:
        :return:
        """
        note_list = []
        try:
            success, msg, all_note_info = self.xhs_apis.get_user_all_notes(user_url, cookies_str, proxies)
            if success:
                logger.info(f'用户 {user_url} 作品数量: {len(all_note_info)}')
                for simple_note_info in all_note_info:
                    note_url = f"https://www.xiaohongshu.com/explore/{simple_note_info['note_id']}?xsec_token={simple_note_info['xsec_token']}"
                    note_list.append(note_url)
            if save_choice == 'all' or save_choice == 'excel':
                excel_name = user_url.split('/')[-1].split('?')[0]
            self.spider_some_note(note_list, cookies_str, base_path, save_choice, excel_name, proxies)
        except Exception as e:
            success = False
            msg = e
        logger.info(f'爬取用户所有视频 {user_url}: {success}, msg: {msg}')
        return note_list, success, msg

    def spider_some_search_note(self, query: str, require_num: int, cookies_str: str, base_path: dict, save_choice: str, sort="general", note_type=0,  excel_name: str = '', proxies=None):
        """
            指定数量搜索笔记，设置排序方式和笔记类型和笔记数量
            :param query 搜索的关键词
            :param require_num 搜索的数量
            :param cookies_str 你的cookies
            :param base_path 保存路径
            :param sort 排序方式 general:综合排序, time_descending:时间排序, popularity_descending:热度排序
            :param note_type 笔记类型 0:全部, 1:视频, 2:图文
            返回搜索的结果
        """
        note_list = []
        try:
            success, msg, notes = self.xhs_apis.search_some_note(query, require_num, cookies_str, sort, note_type, proxies)
            if success:
                notes = list(filter(lambda x: x['model_type'] == "note", notes))
                logger.info(f'搜索关键词 {query} 笔记数量: {len(notes)}')
                for note in notes:
                    note_url = f"https://www.xiaohongshu.com/explore/{note['id']}?xsec_token={note['xsec_token']}"
                    note_list.append(note_url)
            if save_choice == 'all' or save_choice == 'excel':
                excel_name = query
            self.spider_some_note(note_list, cookies_str, base_path, save_choice, excel_name, proxies)
        except Exception as e:
            success = False
            msg = e
        logger.info(f'搜索关键词 {query} 笔记: {success}, msg: {msg}')
        return note_list, success, msg

class DataSpiderGUI(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("小红书采集助手 v1.6")
        self.geometry("800x600")

        # 初始化爬虫实例
        self.spider = Data_Spider()

        # 创建UI组件
        self.create_widgets()

    def create_widgets(self):
        # 顶部配置区域
        config_frame = ttk.LabelFrame(self, text="全局配置")
        config_frame.pack(padx=10, pady=5, fill=tk.X)

        # Cookies配置
        ttk.Label(config_frame, text="Cookies:").grid(row=0, column=0, padx=5)
        self.cookies_entry = ttk.Entry(config_frame, width=70)
        self.cookies_entry.grid(row=0, column=1, padx=5)
        ttk.Button(config_frame, text="导入Cookies", command=self.load_cookies).grid(row=0, column=2)

        # 代理配置
        ttk.Label(config_frame, text="代理(可选):").grid(row=1, column=0, padx=5)
        self.proxy_entry = ttk.Entry(config_frame, width=70)
        self.proxy_entry.grid(row=1, column=1, padx=5)

        # 保存路径配置
        ttk.Label(config_frame, text="保存路径:").grid(row=2, column=0, padx=5)
        self.path_entry = ttk.Entry(config_frame, width=70)
        self.path_entry.grid(row=2, column=1, padx=5)
        ttk.Button(config_frame, text="选择路径", command=self.choose_path).grid(row=2, column=2)

        # 标签页
        self.notebook = ttk.Notebook(self)
        self.notebook.pack(padx=10, pady=5, fill=tk.BOTH, expand=True)

        # 修改各标签页的按钮创建部分
        self.create_note_tab()
        self.create_user_tab()
        self.create_search_tab()

    def create_note_tab(self):
        """笔记链接爬取页"""
        tab = ttk.Frame(self.notebook)
        self.notebook.add(tab, text="笔记爬取")

        # URL输入
        ttk.Label(tab, text="笔记URL（每行一个）:").pack(anchor=tk.W, padx=10, pady=5)
        self.note_urls_text = tk.Text(tab, height=10)
        self.note_urls_text.pack(padx=10, pady=5, fill=tk.X)

        # 保存选项
        save_frame = ttk.Frame(tab)
        save_frame.pack(padx=10, pady=5, fill=tk.X)

        self.note_save_var = tk.StringVar(value="all")
        ttk.Radiobutton(save_frame, text="全部保存", variable=self.note_save_var, value="all").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅媒体", variable=self.note_save_var, value="media").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅Excel", variable=self.note_save_var, value="excel").pack(side=tk.LEFT)

        # Excel名称
        ttk.Label(tab, text="Excel文件名:").pack(anchor=tk.W, padx=10)
        self.note_excel_entry = ttk.Entry(tab)
        self.note_excel_entry.pack(padx=10, pady=5, fill=tk.X)

        # ttk.Button(tab, text="开始爬取", command=self.start_note_spider).pack(pady=10)

        self.note_start_button = ttk.Button(tab, text="开始爬取", command=self.start_note_spider)
        self.note_start_button.pack(pady=10)

    def create_user_tab(self):
        """用户爬取页"""
        tab = ttk.Frame(self.notebook)
        self.notebook.add(tab, text="用户爬取")

        # 用户URL
        ttk.Label(tab, text="用户主页URL:").pack(anchor=tk.W, padx=10, pady=5)
        self.user_url_entry = ttk.Entry(tab)
        self.user_url_entry.pack(padx=10, pady=5, fill=tk.X)

        # 保存选项
        save_frame = ttk.Frame(tab)
        save_frame.pack(padx=10, pady=5, fill=tk.X)

        self.user_save_var = tk.StringVar(value="all")
        ttk.Radiobutton(save_frame, text="全部保存", variable=self.user_save_var, value="all").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅媒体", variable=self.user_save_var, value="media").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅Excel", variable=self.user_save_var, value="excel").pack(side=tk.LEFT)

        # Excel名称
        ttk.Label(tab, text="Excel文件名（留空自动生成）:").pack(anchor=tk.W, padx=10)
        self.user_excel_entry = ttk.Entry(tab)
        self.user_excel_entry.pack(padx=10, pady=5, fill=tk.X)

        # ttk.Button(tab, text="开始爬取", command=self.start_user_spider).pack(pady=10)

        self.user_start_button = ttk.Button(tab, text="开始爬取", command=self.start_user_spider)
        self.user_start_button.pack(pady=10)

    def create_search_tab(self):
        """关键词搜索页"""
        tab = ttk.Frame(self.notebook)
        self.notebook.add(tab, text="关键词搜索")

        # 关键词
        ttk.Label(tab, text="搜索关键词:").pack(anchor=tk.W, padx=10, pady=5)
        self.search_entry = ttk.Entry(tab)
        self.search_entry.pack(padx=10, pady=5, fill=tk.X)

        # 数量
        ttk.Label(tab, text="爬取数量:").pack(anchor=tk.W, padx=10)
        self.num_entry = ttk.Entry(tab)
        self.num_entry.pack(padx=10, pady=5, fill=tk.X)

        # 排序方式
        ttk.Label(tab, text="排序方式:").pack(anchor=tk.W, padx=10)
        self.sort_combo = ttk.Combobox(tab, values=["综合排序", "时间排序", "热度排序"])
        self.sort_combo.current(0)
        self.sort_combo.pack(padx=10, pady=5, fill=tk.X)

        # 笔记类型
        ttk.Label(tab, text="笔记类型:").pack(anchor=tk.W, padx=10)
        self.type_combo = ttk.Combobox(tab, values=["全部", "视频", "图文"])
        self.type_combo.current(0)
        self.type_combo.pack(padx=10, pady=5, fill=tk.X)

        # 保存选项
        save_frame = ttk.Frame(tab)
        save_frame.pack(padx=10, pady=5, fill=tk.X)

        self.search_save_var = tk.StringVar(value="all")
        ttk.Radiobutton(save_frame, text="全部保存", variable=self.search_save_var, value="all").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅媒体", variable=self.search_save_var, value="media").pack(side=tk.LEFT)
        ttk.Radiobutton(save_frame, text="仅Excel", variable=self.search_save_var, value="excel").pack(side=tk.LEFT)

        # Excel名称
        ttk.Label(tab, text="Excel文件名:").pack(anchor=tk.W, padx=10)
        self.search_excel_entry = ttk.Entry(tab)
        self.search_excel_entry.pack(padx=10, pady=5, fill=tk.X)

        # ttk.Button(tab, text="开始爬取", command=self.start_search_spider).pack(pady=10)

        self.search_start_button = ttk.Button(tab, text="开始爬取", command=self.start_search_spider)
        self.search_start_button.pack(pady=10)

    def load_cookies(self):
        """加载Cookies文件"""
        filepath = filedialog.askopenfilename()
        if filepath:
            with open(filepath, 'r') as f:
                self.cookies_entry.delete(0, tk.END)
                self.cookies_entry.insert(0, f.read())

    def choose_path(self):
        """选择保存路径"""
        path = filedialog.askdirectory()
        if path:
            # 使用 os.path.normpath 标准化路径分隔符
            path = os.path.normpath(path)
            self.path_entry.delete(0, tk.END)
            self.path_entry.insert(0, path)

    def get_base_path(self):
        """获取基础保存路径"""
        path = self.path_entry.get()
        if not path:
            messagebox.showerror("错误", "请先选择保存路径")
            return None
        return {
            'media': os.path.join(path, 'media'),
            'excel': os.path.join(path, 'excel')
        }

    def get_proxies(self):
        """获取代理配置"""
        proxy = self.proxy_entry.get().strip()
        return {'http': proxy, 'https': proxy} if proxy else None

    def start_note_spider(self):
        """启动笔记爬取"""
        cookies = self.cookies_entry.get()
        base_path = self.get_base_path()
        if not cookies or not base_path:
            return

        urls = self.note_urls_text.get("1.0", tk.END).strip().split('\n')
        urls = [url.strip() for url in urls if url.strip()]
        if not urls:
            messagebox.showerror("错误", "请输入至少一个笔记URL")
            return

        save_choice = self.note_save_var.get()
        excel_name = self.note_excel_entry.get().strip()
        if save_choice in ('all', 'excel') and not excel_name:
            messagebox.showerror("错误", "Excel文件名不能为空")
            return

        proxies = self.get_proxies()

        # 修改按钮状态
        self.note_start_button.config(text="正在爬取...", state=tk.DISABLED)

        def run():
            try:
                # 创建必要目录
                os.makedirs(base_path['media'], exist_ok=True)
                os.makedirs(base_path['excel'], exist_ok=True)

                self.spider.spider_some_note(
                    urls, cookies, base_path, save_choice, excel_name, proxies
                )

                self.after(0, lambda: messagebox.showinfo("成功", "笔记爬取完成！"))

                # messagebox.showinfo("成功", "笔记爬取完成！")
            except Exception as e:
                self.after(0, lambda: messagebox.showerror("错误", str(e)))

            finally:
                self.after(0, lambda: self.note_start_button.config(
                    text="开始爬取",
                    state=tk.NORMAL
                ))

        threading.Thread(target=run, daemon=True).start()

    def start_user_spider(self):
        """启动用户爬取"""
        cookies = self.cookies_entry.get()
        base_path = self.get_base_path()
        if not cookies or not base_path:
            return

        user_url = self.user_url_entry.get().strip()
        if not user_url:
            messagebox.showerror("错误", "请输入用户主页URL")
            return

        save_choice = self.user_save_var.get()
        excel_name = self.user_excel_entry.get().strip()

        proxies = self.get_proxies()
        # 修改按钮状态
        self.user_start_button.config(text="正在爬取...", state=tk.DISABLED)

        def run():
            try:
                # 创建必要目录
                os.makedirs(base_path['media'], exist_ok=True)
                os.makedirs(base_path['excel'], exist_ok=True)

                self.spider.spider_user_all_note(
                    user_url, cookies, base_path, save_choice, excel_name, proxies
                )
                # messagebox.showinfo("成功", "用户笔记爬取完成！")
                self.after(0, lambda: messagebox.showinfo("成功", "用户笔记爬取完成！"))

            except Exception as e:
                self.after(0, lambda: messagebox.showerror("错误", str(e)))
                # messagebox.showerror("错误", str(e))

            finally:
                self.after(0, lambda: self.user_start_button.config(
                    text="开始爬取",
                    state=tk.NORMAL
                ))

        threading.Thread(target=run).start()

    def start_search_spider(self):
        """启动关键词搜索爬取"""
        cookies = self.cookies_entry.get()
        base_path = self.get_base_path()
        if not cookies or not base_path:
            return

        query = self.search_entry.get().strip()
        if not query:
            messagebox.showerror("错误", "请输入搜索关键词")
            return

        try:
            num = int(self.num_entry.get())
        except ValueError:
            messagebox.showerror("错误", "请输入有效的数量")
            return

        sort_map = {"综合排序": "general", "时间排序": "time_descending", "热度排序": "popularity_descending"}
        sort = sort_map[self.sort_combo.get()]

        note_type = self.type_combo.current()  # 0:全部, 1:视频, 2:图文

        save_choice = self.search_save_var.get()
        excel_name = self.search_excel_entry.get().strip()
        if save_choice in ('all', 'excel') and not excel_name:
            messagebox.showerror("错误", "Excel文件名不能为空")
            return

        proxies = self.get_proxies()

        self.search_start_button.config(text="正在爬取...", state=tk.DISABLED)

        def run():
            try:
                # 创建必要目录
                os.makedirs(base_path['media'], exist_ok=True)
                os.makedirs(base_path['excel'], exist_ok=True)

                self.spider.spider_some_search_note(
                    query, num, cookies, base_path, save_choice, sort, note_type, excel_name, proxies
                )
                # messagebox.showinfo("成功", "关键词搜索爬取完成！")
                self.after(0, lambda: messagebox.showinfo("成功", "关键词搜索爬取完成！"))
            except Exception as e:
                self.after(0, lambda: messagebox.showerror("错误", str(e)))
                # messagebox.showerror("错误", str(e))
            finally:
                self.after(0, lambda: self.search_start_button.config(
                    text="开始爬取",
                    state=tk.NORMAL
                ))


        threading.Thread(target=run).start()


if __name__ == "__main__":
    app = DataSpiderGUI()
    app.mainloop()